package datastructure.trie;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;


class 多次搜索_面试题17_重做版 {
    public int[][] multiSearch(String big, String[] smalls) {
        Trie trie = new Trie(smalls);
        Map<Integer, List<Integer>> map = new HashMap<>();
        for (int i = 0; i < big.length(); i++) {
            String seq = big.substring(i);
            List<Integer> poss = trie.search(seq);
            for (Integer pos : poss) {
                List<Integer> list = map.computeIfAbsent(pos, k -> new ArrayList<>());
                list.add(i);
            }
        }
        int[][] res = new int[smalls.length][];
        for (int i = 0; i < res.length; i++) {
            List<Integer> p = map.get(i);
            if (p == null) {
                res[i] = new int[]{};
            } else {
                int[] pi = new int[p.size()];
                int idx = 0;
                for (Integer in : p) {
                    pi[idx++] = in;
                }
                res[i] = pi;
            }

        }
        return res;

    }

    class Trie {
        TrieNode root = new TrieNode();

        public Trie(String[] smalls) {
            for (int s = 0; s < smalls.length; s++) {
                TrieNode node = root;
                for (int i = 0; i < smalls[s].length(); i++) {
                    int c = smalls[s].charAt(i) - 'a';
                    if (node.child[c] == null) node.child[c] = new TrieNode();
                    node = node.child[c];
                }
                node.idx = s;
            }
        }

        // 片段中搜索哪些字符串
        public List<Integer> search(String sub) {
            TrieNode node = root;
            List<Integer> res = new ArrayList<>();
            for (int i = 0; i < sub.length(); i++) {
                int c = sub.charAt(i) - 'a';
                if (node.child[c] == null) break;
                node = node.child[c];
                if (node.idx != -1) res.add(node.idx);
            }
            return res;
        }
    }

    class TrieNode {
        TrieNode[] child = new TrieNode[26];
        Integer idx = -1;  // 保存smalls 中的下标 直接回头匹配

        public TrieNode() {

        }
    }
}

/**
 * 转化为对每个子字符串   前缀有哪些的前缀树问题
 * https://leetcode-cn.com/problems/multi-search-lcci/
 */
public class 多次搜索_17 {

    /**
     * 对每个子序列   查看是否存在前缀
     */
    public int[][] multiSearch(String big, String[] smalls) {
        Trie trie = new Trie(smalls);
        Map<String, List<Integer>> map = new HashMap<>();
        for (int i = 0; i < big.length(); i++) {
            String subText = big.substring(i);
            final List<String> list = trie.searchSub(subText);
            for (String smallItem : list) {
                List<Integer> count = map.computeIfAbsent(smallItem, k -> new ArrayList<>());
                count.add(i);
            }
        }
        int[][] res = new int[smalls.length][];
        for (int i = 0; i < smalls.length; i++) {
            final List<Integer> list = map.get(smalls[i]);
            int[] col;
            if (list == null || list.size() == 0) {
                col = new int[]{};
            } else {
                col = new int[list.size()];
                int idx = 0;
                for (int colItem : list) {
                    col[idx++] = colItem;
                }
            }

            res[i] = col;
        }
        return res;
    }


    private class Trie {
        TrieNode root = new TrieNode();

        public Trie(String[] smalls) {
            for (String small : smalls) {
                TrieNode node = root;
                for (int i = 0; i < small.length(); i++) {
                    int idx = small.charAt(i) - 'a';
                    if (node.childs[idx] == null) {
                        node.childs[idx] = new TrieNode();
                    }
                    node = node.childs[idx];
                }
                node.sub = small;
            }
        }

        /**
         * 该text有哪些前缀存在 trie中
         */
        public List<String> searchSub(String text) {
            List<String> res = new ArrayList<>();
            TrieNode node = root;
            for (int i = 0; i < text.length(); i++) {
                int idx = text.charAt(i) - 'a';
                if (node.childs[idx] == null) {
                    break;
                }
                // 有数组中的子串是text的前缀
                node = node.childs[idx];
                if (node.sub != null) {
                    res.add(node.sub);
                }
            }
            return res;
        }
    }

    private class TrieNode {
        TrieNode[] childs = new TrieNode[26];
        String sub = null;
    }

}
